# Copyright 2018-2021 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This module contains basis set parameters defining Gaussian-type orbitals for a selected number of
atoms. The data are taken from the Basis Set Exchange `library <https://www.basissetexchange.org>`_.
The current data includes the STO-3G, 6-31G, 6-311G and CC-PVDZ, basis sets for elements with atomic
numbers 1-10.
"""

import itertools

# Note: Below markers are added to prevent reformatting of this dictionary by black
# fmt: off
# IUPAC Periodic Table of the Elements: https://iupac.org/what-we-do/periodic-table-of-elements/
atomic_numbers = {
    'H': 1, 'He': 2, # Period 1
    'Li': 3, 'Be': 4, 'B': 5, 'C': 6, 'N': 7, 'O': 8, 'F': 9, 'Ne': 10, # Period 2
    'Na': 11, 'Mg': 12, 'Al': 13, 'Si': 14, 'P': 15, 'S': 16, 'Cl': 17, 'Ar': 18, # Period 3
    'K': 19, 'Ca': 20, 'Sc': 21, 'Ti': 22, 'V': 23, 'Cr': 24, 'Mn': 25, 'Fe': 26, 'Co': 27, # Period 4
    'Ni': 28, 'Cu': 29, 'Zn': 30, 'Ga': 31, 'Ge': 32, 'As': 33, 'Se': 34, 'Br': 35, 'Kr': 36,
    'Rb': 37, 'Sr': 38, 'Y': 39, 'Zr': 40, 'Nb': 41, 'Mo': 42, 'Tc': 43, 'Ru': 44, 'Rh': 45, # Period 5
    'Pd': 46, 'Ag': 47, 'Cd': 48, 'In': 49, 'Sn': 50, 'Sb': 51, 'Te': 52, 'I': 53, 'Xe': 54,
    'Cs': 55, 'Ba': 56, # Period 6
    'La': 57, 'Ce': 58, 'Pr': 59, 'Nd': 60, 'Pm': 61, 'Sm': 62, 'Eu': 63, 'Gd': 64, 'Tb': 65, # Lanthanides
    'Dy': 66, 'Ho': 67, 'Er': 68, 'Tm': 69, 'Yb': 70, 'Lu': 71,
    'Hf': 72, 'Ta': 73, 'W': 74, 'Re': 75, 'Os': 76, 'Ir': 77, 'Pt': 78, 'Au': 79, 'Hg': 80, # Period 6
    'Tl': 81, 'Pb': 82, 'Bi': 83, 'Po': 84, 'At': 85, 'Rn': 86,
    'Fr': 87, 'Ra': 88, # Period 7
    'Ac': 89, 'Th': 90, 'Pa': 91, 'U': 92, 'Np': 93, 'Pu': 94, 'Am': 95, 'Cm': 96, 'Bk': 97, # Actinides
    'Cf': 98, 'Es': 99, 'Fm': 100, 'Md': 101, 'No': 102, 'Lr': 103,
    'Rf': 104, 'Db': 105, 'Sg': 106, 'Bh': 107, 'Hs': 108, 'Mt': 109, 'Ds': 110, 'Rg': 111, # Period 7
    'Cn': 112, 'Nh': 113, 'Fl': 114, 'Mc': 115, 'Lv': 116, 'Ts': 117, 'Og': 118
}
# fmt: on


STO3G = {
    "H": {
        "orbitals": ["S"],
        "exponents": [[0.3425250914e01, 0.6239137298e00, 0.1688554040e00]],
        "coefficients": [[0.1543289673e00, 0.5353281423e00, 0.4446345422e00]],
    },
    "He": {
        "orbitals": ["S"],
        "exponents": [[0.6362421394e01, 0.1158922999e01, 0.3136497915e00]],
        "coefficients": [[0.1543289673e00, 0.5353281423e00, 0.4446345422e00]],
    },
    "Li": {
        "orbitals": ["S", "SP"],
        "exponents": [
            [0.1611957475e02, 0.2936200663e01, 0.7946504870e00],
            [0.6362897469e00, 0.1478600533e00, 0.4808867840e-01],
        ],
        "coefficients": [
            [0.1543289673e00, 0.5353281423e00, 0.4446345422e00],
            [-0.9996722919e-01, 0.3995128261e00, 0.7001154689e00],
            [0.1559162750e00, 0.6076837186e00, 0.3919573931e00],
        ],
    },
    "Be": {
        "orbitals": ["S", "SP"],
        "exponents": [
            [0.3016787069e02, 0.5495115306e01, 0.1487192653e01],
            [0.1314833110e01, 0.3055389383e00, 0.9937074560e-01],
        ],
        "coefficients": [
            [0.1543289673e00, 0.5353281423e00, 0.4446345422e00],
            [-0.9996722919e-01, 0.3995128261e00, 0.7001154689e00],
            [0.1559162750e00, 0.6076837186e00, 0.3919573931e00],
        ],
    },
    "B": {
        "orbitals": ["S", "SP"],
        "exponents": [
            [0.4879111318e02, 0.8887362172e01, 0.2405267040e01],
            [0.2236956142e01, 0.5198204999e00, 0.1690617600e00],
        ],
        "coefficients": [
            [0.1543289673e00, 0.5353281423e00, 0.4446345422e00],
            [-0.9996722919e-01, 0.3995128261e00, 0.7001154689e00],
            [0.1559162750e00, 0.6076837186e00, 0.3919573931e00],
        ],
    },
    "C": {
        "orbitals": ["S", "SP"],
        "exponents": [
            [0.7161683735e02, 0.1304509632e02, 0.3530512160e01],
            [0.2941249355e01, 0.6834830964e00, 0.2222899159e00],
        ],
        "coefficients": [
            [0.1543289673e00, 0.5353281423e00, 0.4446345422e00],
            [-0.9996722919e-01, 0.3995128261e00, 0.7001154689e00],
            [0.1559162750e00, 0.6076837186e00, 0.3919573931e00],
        ],
    },
    "N": {
        "orbitals": ["S", "SP"],
        "exponents": [
            [0.9910616896e02, 0.1805231239e02, 0.4885660238e01],
            [0.3780455879e01, 0.8784966449e00, 0.2857143744e00],
        ],
        "coefficients": [
            [0.1543289673e00, 0.5353281423e00, 0.4446345422e00],
            [-0.9996722919e-01, 0.3995128261e00, 0.7001154689e00],
            [0.1559162750e00, 0.6076837186e00, 0.3919573931e00],
        ],
    },
    "O": {
        "orbitals": ["S", "SP"],
        "exponents": [
            [0.1307093214e03, 0.2380886605e02, 0.6443608313e01],
            [0.5033151319e01, 0.1169596125e01, 0.3803889600e00],
        ],
        "coefficients": [
            [0.1543289673e00, 0.5353281423e00, 0.4446345422e00],
            [-0.9996722919e-01, 0.3995128261e00, 0.7001154689e00],
            [0.1559162750e00, 0.6076837186e00, 0.3919573931e00],
        ],
    },
    "F": {
        "orbitals": ["S", "SP"],
        "exponents": [
            [0.1666791340e03, 0.3036081233e02, 0.8216820672e01],
            [0.6464803249e01, 0.1502281245e01, 0.4885884864e00],
        ],
        "coefficients": [
            [0.1543289673e00, 0.5353281423e00, 0.4446345422e00],
            [-0.9996722919e-01, 0.3995128261e00, 0.7001154689e00],
            [0.1559162750e00, 0.6076837186e00, 0.3919573931e00],
        ],
    },
    "Ne": {
        "orbitals": ["S", "SP"],
        "exponents": [
            [0.2070156070e03, 0.3770815124e02, 0.1020529731e02],
            [0.8246315120e01, 0.1916266291e01, 0.6232292721e00],
        ],
        "coefficients": [
            [0.1543289673e00, 0.5353281423e00, 0.4446345422e00],
            [-0.9996722919e-01, 0.3995128261e00, 0.7001154689e00],
            [0.1559162750e00, 0.6076837186e00, 0.3919573931e00],
        ],
    },
}

POPLE631G = {
    "H": {
        "orbitals": ["S", "S"],
        "exponents": [[18.73113696, 2.825394365, 0.6401216923], [0.1612777588]],
        "coefficients": [[0.03349460434, 0.2347269535, 0.8137573261], [1.0]],
    },
    "He": {
        "orbitals": ["S", "S"],
        "exponents": [[38.421634, 5.77803, 1.241774], [0.297964]],
        "coefficients": [[0.04013973935, 0.261246097, 0.7931846246], [1.0]],
    },
    "Li": {
        "orbitals": ["S", "SP", "SP"],
        "exponents": [
            [642.418915, 96.7985153, 22.0911212, 6.20107025, 1.93511768, 0.636735789],
            [2.324918408, 0.6324303556, 0.07905343475],
            [0.03596197175],
        ],
        "coefficients": [
            [0.00214260781, 0.0162088715, 0.0773155725, 0.245786052, 0.470189004, 0.345470845],
            [-0.03509174574, -0.1912328431, 1.083987795],
            [0.008941508043, 0.141009464, 0.9453636953],
            [1.0],
            [1.0],
        ],
    },
    "Be": {
        "orbitals": ["S", "SP", "SP"],
        "exponents": [
            [1264.58569, 189.936806, 43.159089, 12.0986627, 3.80632322, 1.2728903],
            [3.196463098, 0.7478133038, 0.2199663302],
            [0.0823099007],
        ],
        "coefficients": [
            [0.00194475759, 0.014835052, 0.07209054629, 0.23715415, 0.4691986519, 0.3565202279],
            [-0.1126487285, -0.2295064079, 1.186916764],
            [0.0559801998, 0.261550611, 0.7939723389],
            [1.0],
            [1.0],
        ],
    },
    "B": {
        "orbitals": ["S", "SP", "SP"],
        "exponents": [
            [2068.88225, 310.64957, 70.683033, 19.8610803, 6.29930484, 2.12702697],
            [4.727971071, 1.190337736, 0.3594116829],
            [0.1267512469],
        ],
        "coefficients": [
            [0.00186627459, 0.0142514817, 0.0695516185, 0.232572933, 0.467078712, 0.36343144],
            [-0.1303937974, -0.1307889514, 1.130944484],
            [0.07459757992, 0.3078466771, 0.7434568342],
            [1.0],
            [1.0],
        ],
    },
    "C": {
        "orbitals": ["S", "SP", "SP"],
        "exponents": [
            [3047.52488, 457.369518, 103.948685, 29.2101553, 9.28666296, 3.16392696],
            [7.86827235, 1.88128854, 0.544249258],
            [0.1687144782],
        ],
        "coefficients": [
            [
                0.001834737132,
                0.01403732281,
                0.06884262226,
                0.2321844432,
                0.4679413484,
                0.3623119853,
            ],
            [-0.1193324198, -0.1608541517, 1.143456438],
            [0.06899906659, 0.316423961, 0.7443082909],
            [1.0],
            [1.0],
        ],
    },
    "N": {
        "orbitals": ["S", "SP", "SP"],
        "exponents": [
            [4173.51146, 627.457911, 142.902093, 40.2343293, 12.8202129, 4.39043701],
            [11.62636186, 2.716279807, 0.7722183966],
            [0.2120314975],
        ],
        "coefficients": [
            [0.00183477216, 0.013994627, 0.06858655181, 0.232240873, 0.4690699481, 0.3604551991],
            [-0.1149611817, -0.1691174786, 1.145851947],
            [0.06757974388, 0.3239072959, 0.7408951398],
            [1.0],
            [1.0],
        ],
    },
    "O": {
        "orbitals": ["S", "SP", "SP"],
        "exponents": [
            [5484.67166, 825.234946, 188.046958, 52.9645, 16.8975704, 5.79963534],
            [15.53961625, 3.599933586, 1.01376175],
            [0.2700058226],
        ],
        "coefficients": [
            [0.00183107443, 0.0139501722, 0.0684450781, 0.232714336, 0.470192898, 0.358520853],
            [-0.1107775495, -0.1480262627, 1.130767015],
            [0.07087426823, 0.3397528391, 0.7271585773],
            [1.0],
            [1.0],
        ],
    },
    "F": {
        "orbitals": ["S", "SP", "SP"],
        "exponents": [
            [7001.71309, 1051.36609, 239.28569, 67.3974453, 21.5199573, 7.4031013],
            [20.8479528, 4.80830834, 1.34406986],
            [0.358151393],
        ],
        "coefficients": [
            [
                0.001819616901,
                0.01391607961,
                0.06840532453,
                0.2331857601,
                0.4712674392,
                0.3566185462,
            ],
            [-0.1085069751, -0.1464516581, 1.128688581],
            [0.07162872424, 0.3459121027, 0.7224699564],
            [1.0],
            [1.0],
        ],
    },
    "Ne": {
        "orbitals": ["S", "SP", "SP"],
        "exponents": [
            [8425.85153, 1268.5194, 289.621414, 81.859004, 26.2515079, 9.09472051],
            [26.532131, 6.10175501, 1.69627153],
            [0.4458187],
        ],
        "coefficients": [
            [0.00188434805, 0.0143368994, 0.07010962331, 0.237373266, 0.4730071261, 0.348401241],
            [-0.1071182872, -0.1461638213, 1.127773503],
            [0.07190958851, 0.349513372, 0.7199405121],
            [1.0],
            [1.0],
        ],
    },
}

POPLE6311G = {
    "H": {
        "orbitals": ["S", "S", "S"],
        "exponents": [[33.865, 5.09479, 1.15879], [0.32584], [0.102741]],
        "coefficients": [[0.0254938, 0.190373, 0.852161], [1.0], [1.0]],
    },
    "He": {
        "orbitals": ["S", "S", "S"],
        "exponents": [[98.1243, 14.7689, 3.31883], [0.874047], [0.244564]],
        "coefficients": [[0.0287452, 0.208061, 0.837635], [1.0], [1.0]],
    },
    "Li": {
        "orbitals": ["S", "SP", "SP", "SP"],
        "exponents": [
            [900.46, 134.433, 30.4365, 8.62639, 2.48332, 0.303179],
            [4.8689, 0.856924, 0.243227],
            [0.063507],
            [0.0243683],
        ],
        "coefficients": [
            [0.00228704, 0.017635, 0.0873434, 0.280977, 0.658741, 0.118712],
            [0.0933293, 0.943045, -0.00279827],
            [0.0327661, 0.159792, 0.885667],
            [1.0],
            [1.0],
            [1.0],
            [1.0],
        ],
    },
    "Be": {
        "orbitals": ["S", "SP", "SP", "SP"],
        "exponents": [
            [1682.8, 251.715, 57.4116, 16.5171, 4.85364, 0.626863],
            [8.30938, 1.74075, 0.485816],
            [0.163613],
            [0.0567285],
        ],
        "coefficients": [
            [0.00228574, 0.0175938, 0.0863315, 0.281835, 0.640594, 0.144467],
            [0.108621, 0.927301, -0.00297169],
            [0.0361344, 0.216958, 0.841839],
            [1.0],
            [1.0],
            [1.0],
            [1.0],
        ],
    },
    "B": {
        "orbitals": ["S", "SP", "SP", "SP"],
        "exponents": [
            [2858.89, 428.14, 97.5282, 27.9693, 8.21577, 1.11278],
            [13.2415, 3.00166, 0.912856],
            [0.315454],
            [0.0988563],
        ],
        "coefficients": [
            [0.00215375, 0.0165823, 0.082187, 0.276618, 0.629316, 0.17377],
            [0.117443, 0.918002, -0.00265105],
            [0.04181, 0.236575, 0.816214],
            [1.0],
            [1.0],
            [1.0],
            [1.0],
        ],
    },
    "C": {
        "orbitals": ["S", "SP", "SP", "SP"],
        "exponents": [
            [4563.24, 682.024, 154.973, 44.4553, 13.029, 1.82773],
            [20.9642, 4.80331, 1.45933],
            [0.483456],
            [0.145585],
        ],
        "coefficients": [
            [0.00196665, 0.0152306, 0.0761269, 0.260801, 0.616462, 0.221006],
            [0.11466, 0.919999, -0.00303068],
            [0.0402487, 0.237594, 0.815854],
            [1.0],
            [1.0],
            [1.0],
            [1.0],
        ],
    },
    "N": {
        "orbitals": ["S", "SP", "SP", "SP"],
        "exponents": [
            [6293.48, 949.044, 218.776, 63.6916, 18.8282, 2.72023],
            [30.6331, 7.02614, 2.11205],
            [0.684009],
            [0.200878],
        ],
        "coefficients": [
            [0.00196979, 0.0149613, 0.0735006, 0.248937, 0.60246, 0.256202],
            [0.111906, 0.921666, -0.00256919],
            [0.0383119, 0.237403, 0.817592],
            [1.0],
            [1.0],
            [1.0],
            [1.0],
        ],
    },
    "O": {
        "orbitals": ["S", "SP", "SP", "SP"],
        "exponents": [
            [8588.5, 1297.23, 299.296, 87.3771, 25.6789, 3.74004],
            [42.1175, 9.62837, 2.85332],
            [0.905661],
            [0.255611],
        ],
        "coefficients": [
            [0.00189515, 0.0143859, 0.070732, 0.240001, 0.594797, 0.280802],
            [0.113889, 0.920811, -0.00327447],
            [0.0365114, 0.237153, 0.819702],
            [1.0],
            [1.0],
            [1.0],
            [1.0],
        ],
    },
    "F": {
        "orbitals": ["S", "SP", "SP", "SP"],
        "exponents": [
            [11427.1, 1722.35, 395.746, 115.139, 33.6026, 4.91901],
            [55.4441, 12.6323, 3.71756],
            [1.16545],
            [0.321892],
        ],
        "coefficients": [
            [0.00180093, 0.0137419, 0.0681334, 0.233325, 0.589086, 0.299505],
            [0.114536, 0.920512, -0.00337804],
            [0.0354609, 0.237451, 0.820458],
            [1.0],
            [1.0],
            [1.0],
            [1.0],
        ],
    },
    "Ne": {
        "orbitals": ["S", "SP", "SP", "SP"],
        "exponents": [
            [13995.7, 2117.1, 490.425, 143.833, 41.9265, 6.15684],
            [69.1211, 15.835, 4.67326],
            [1.45756],
            [0.397057],
        ],
        "coefficients": [
            [0.00183276, 0.0138827, 0.0680687, 0.231328, 0.58589, 0.305883],
            [0.119149, 0.917375, -0.00405839],
            [0.0356574, 0.239477, 0.818461],
            [1.0],
            [1.0],
            [1.0],
            [1.0],
        ],
    },
}

CCPVDZ = {
    "H": {
        "orbitals": ["S", "S", "P"],
        "exponents": [[13.01, 1.962, 0.4446, 0.122], [13.01, 1.962, 0.4446, 0.122], [0.727]],
        "coefficients": [[0.019685, 0.137977, 0.478148, 0.50124], [0.0, 0.0, 0.0, 1.0], [1.0]],
    },
    "He": {
        "orbitals": ["S", "S", "P"],
        "exponents": [[38.36, 5.77, 1.24, 0.2976], [38.36, 5.77, 1.24, 0.2976], [1.275]],
        "coefficients": [[0.023809, 0.154891, 0.469987, 0.513027], [0.0, 0.0, 0.0, 1.0], [1.0]],
    },
    "Li": {
        "orbitals": ["S", "S", "S", "P", "P", "D"],
        "exponents": [
            [1469.0, 220.5, 50.26, 14.24, 4.581, 1.58, 0.564, 0.07345, 0.02805],
            [1469.0, 220.5, 50.26, 14.24, 4.581, 1.58, 0.564, 0.07345, 0.02805],
            [1469.0, 220.5, 50.26, 14.24, 4.581, 1.58, 0.564, 0.07345, 0.02805],
            [1.534, 0.2749, 0.07362, 0.02403],
            [1.534, 0.2749, 0.07362, 0.02403],
            [0.1144],
        ],
        "coefficients": [
            [
                0.000766,
                0.005892,
                0.029671,
                0.10918,
                0.282789,
                0.453123,
                0.274774,
                0.009751,
                -0.00318,
            ],
            [
                -0.00012,
                -0.000923,
                -0.004689,
                -0.017682,
                -0.048902,
                -0.096009,
                -0.13638,
                0.575102,
                0.517661,
            ],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
            [0.022784, 0.139107, 0.500375, 0.508474],
            [0.0, 0.0, 0.0, 1.0],
            [1.0],
        ],
    },
    "Be": {
        "orbitals": ["S", "S", "S", "P", "P", "D"],
        "exponents": [
            [2940.0, 441.2, 100.5, 28.43, 9.169, 3.196, 1.159, 0.1811, 0.0589],
            [2940.0, 441.2, 100.5, 28.43, 9.169, 3.196, 1.159, 0.1811, 0.0589],
            [2940.0, 441.2, 100.5, 28.43, 9.169, 3.196, 1.159, 0.1811, 0.0589],
            [3.619, 0.711, 0.1951, 0.06018],
            [3.619, 0.711, 0.1951, 0.06018],
            [0.2354],
        ],
        "coefficients": [
            [
                0.00068,
                0.005236,
                0.026606,
                0.099993,
                0.269702,
                0.451469,
                0.295074,
                0.012587,
                -0.003756,
            ],
            [
                -0.000123,
                -0.000966,
                -0.004831,
                -0.019314,
                -0.05328,
                -0.120723,
                -0.133435,
                0.530767,
                0.580117,
            ],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
            [0.029111, 0.169365, 0.513458, 0.479338],
            [0.0, 0.0, 0.0, 1.0],
            [1.0],
        ],
    },
    "B": {
        "orbitals": ["S", "S", "S", "P", "P", "D"],
        "exponents": [
            [4570.0, 685.9, 156.5, 44.47, 14.48, 5.131, 1.898, 0.3329, 0.1043],
            [4570.0, 685.9, 156.5, 44.47, 14.48, 5.131, 1.898, 0.3329, 0.1043],
            [4570.0, 685.9, 156.5, 44.47, 14.48, 5.131, 1.898, 0.3329, 0.1043],
            [6.001, 1.241, 0.3364, 0.09538],
            [6.001, 1.241, 0.3364, 0.09538],
            [0.343],
        ],
        "coefficients": [
            [
                0.000696,
                0.005353,
                0.027134,
                0.10138,
                0.272055,
                0.448403,
                0.290123,
                0.014322,
                -0.003486,
            ],
            [
                -0.000139,
                -0.001097,
                -0.005444,
                -0.021916,
                -0.059751,
                -0.138732,
                -0.131482,
                0.539526,
                0.580774,
            ],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
            [0.035481, 0.198072, 0.50523, 0.479499],
            [0.0, 0.0, 0.0, 1.0],
            [1.0],
        ],
    },
    "C": {
        "orbitals": ["S", "S", "S", "P", "P", "D"],
        "exponents": [
            [6665.0, 1000.0, 228.0, 64.71, 21.06, 7.495, 2.797, 0.5215, 0.1596],
            [6665.0, 1000.0, 228.0, 64.71, 21.06, 7.495, 2.797, 0.5215, 0.1596],
            [6665.0, 1000.0, 228.0, 64.71, 21.06, 7.495, 2.797, 0.5215, 0.1596],
            [9.439, 2.002, 0.5456, 0.1517],
            [9.439, 2.002, 0.5456, 0.1517],
            [0.55],
        ],
        "coefficients": [
            [
                0.000692,
                0.005329,
                0.027077,
                0.101718,
                0.27474,
                0.448564,
                0.285074,
                0.015204,
                -0.003191,
            ],
            [
                -0.000146,
                -0.001154,
                -0.005725,
                -0.023312,
                -0.063955,
                -0.149981,
                -0.127262,
                0.544529,
                0.580496,
            ],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
            [0.038109, 0.20948, 0.508557, 0.468842],
            [0.0, 0.0, 0.0, 1.0],
            [1.0],
        ],
    },
    "N": {
        "orbitals": ["S", "S", "S", "P", "P", "D"],
        "exponents": [
            [9046.0, 1357.0, 309.3, 87.73, 28.56, 10.21, 3.838, 0.7466, 0.2248],
            [9046.0, 1357.0, 309.3, 87.73, 28.56, 10.21, 3.838, 0.7466, 0.2248],
            [9046.0, 1357.0, 309.3, 87.73, 28.56, 10.21, 3.838, 0.7466, 0.2248],
            [13.55, 2.917, 0.7973, 0.2185],
            [13.55, 2.917, 0.7973, 0.2185],
            [0.817],
        ],
        "coefficients": [
            [0.0007, 0.005389, 0.027406, 0.103207, 0.278723, 0.44854, 0.278238, 0.01544, -0.002864],
            [
                -0.000153,
                -0.001208,
                -0.005992,
                -0.024544,
                -0.067459,
                -0.158078,
                -0.121831,
                0.549003,
                0.578815,
            ],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
            [0.039919, 0.217169, 0.510319, 0.462214],
            [0.0, 0.0, 0.0, 1.0],
            [1.0],
        ],
    },
    "O": {
        "orbitals": ["S", "S", "S", "P", "P", "D"],
        "exponents": [
            [11720.0, 1759.0, 400.8, 113.7, 37.03, 13.27, 5.025, 1.013, 0.3023],
            [11720.0, 1759.0, 400.8, 113.7, 37.03, 13.27, 5.025, 1.013, 0.3023],
            [11720.0, 1759.0, 400.8, 113.7, 37.03, 13.27, 5.025, 1.013, 0.3023],
            [17.7, 3.854, 1.046, 0.2753],
            [17.7, 3.854, 1.046, 0.2753],
            [1.185],
        ],
        "coefficients": [
            [0.00071, 0.00547, 0.027837, 0.1048, 0.283062, 0.448719, 0.270952, 0.015458, -0.002585],
            [
                -0.00016,
                -0.001263,
                -0.006267,
                -0.025716,
                -0.070924,
                -0.165411,
                -0.116955,
                0.557368,
                0.572759,
            ],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
            [0.043018, 0.228913, 0.508728, 0.460531],
            [0.0, 0.0, 0.0, 1.0],
            [1.0],
        ],
    },
    "F": {
        "orbitals": ["S", "S", "S", "P", "P", "D"],
        "exponents": [
            [14710.0, 2207.0, 502.8, 142.6, 46.47, 16.7, 6.356, 1.316, 0.3897],
            [14710.0, 2207.0, 502.8, 142.6, 46.47, 16.7, 6.356, 1.316, 0.3897],
            [14710.0, 2207.0, 502.8, 142.6, 46.47, 16.7, 6.356, 1.316, 0.3897],
            [22.67, 4.977, 1.347, 0.3471],
            [22.67, 4.977, 1.347, 0.3471],
            [1.64],
        ],
        "coefficients": [
            [
                0.000721,
                0.005553,
                0.028267,
                0.106444,
                0.286814,
                0.448641,
                0.264761,
                0.015333,
                -0.002332,
            ],
            [
                -0.000165,
                -0.001308,
                -0.006495,
                -0.026691,
                -0.07369,
                -0.170776,
                -0.112327,
                0.562814,
                0.568778,
            ],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
            [0.044878, 0.235718, 0.508521, 0.45812],
            [0.0, 0.0, 0.0, 1.0],
            [1.0],
        ],
    },
    "Ne": {
        "orbitals": ["S", "S", "S", "P", "P", "D"],
        "exponents": [
            [17880.0, 2683.0, 611.5, 173.5, 56.64, 20.42, 7.81, 1.653, 0.4869],
            [17880.0, 2683.0, 611.5, 173.5, 56.64, 20.42, 7.81, 1.653, 0.4869],
            [17880.0, 2683.0, 611.5, 173.5, 56.64, 20.42, 7.81, 1.653, 0.4869],
            [28.39, 6.27, 1.695, 0.4317],
            [28.39, 6.27, 1.695, 0.4317],
            [2.202],
        ],
        "coefficients": [
            [
                0.000738,
                0.005677,
                0.028883,
                0.10854,
                0.290907,
                0.448324,
                0.258026,
                0.015063,
                -0.0021,
            ],
            [
                -0.000172,
                -0.001357,
                -0.006737,
                -0.027663,
                -0.076208,
                -0.175227,
                -0.107038,
                0.56705,
                0.565216,
            ],
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
            [0.046087, 0.240181, 0.508744, 0.45566],
            [0.0, 0.0, 0.0, 1.0],
            [1.0],
        ],
    },
}

basis_sets = {
    "sto-3g": STO3G,
    "6-31g": POPLE631G,
    "6-311g": POPLE6311G,
    "cc-pvdz": CCPVDZ,
}


def load_basisset(basis, element):
    r"""Extracts basis set data from the Basis Set Exchange library.

    Args:
        basis (str): name of the basis set
        element (str): atomic symbol of the chemical element

    Returns:
        dict[str, list]: dictionary containing orbital names, and the exponents and contraction
        coefficients of a basis function

    **Example**

    >>> basis = '6-31g'
    >>> element = 'He'
    >>> basis = qml.qchem.load_basisset(basis, element)
    >>> basis
    {'orbitals': ['S', 'S'],
     'exponents': [[38.421634, 5.77803, 1.241774], [0.297964]],
     'coefficients': [[0.04013973935, 0.261246097, 0.7931846246], [1.0]]}
    """
    try:
        # pylint: disable=import-outside-toplevel
        import basis_set_exchange as bse
    except ImportError as Error:
        raise ImportError(
            "This feature requires basis_set_exchange. "
            "It can be installed with: pip install basis-set-exchange."
        ) from Error

    orbital_map = {
        "[0]": "S",
        "[0, 1]": "SP",
        "[0, 1, 2]": "SPD",
        "[1]": "P",
        "[2]": "D",
        "[3]": "F",
        "[4]": "G",
        "[5]": "H",
        "[6]": "I",
    }

    atomic_number = atomic_numbers.get(element, None)
    if atomic_number is None:
        raise ValueError(f"Requested element {element} doesn't exist in the periodic table.")

    data = bse.get_basis(basis)["elements"][str(atomic_number)]["electron_shells"]

    orbitals = []
    exponents = []
    coefficients = []

    for term in data:
        if orbital_map[str(term["angular_momentum"])] == "SP":
            orbitals.append(["S", "P"])
        else:
            orbitals.append(
                [orbital_map[str(term["angular_momentum"])]] * len(term["coefficients"])
            )
        exponents.append(
            [list(map(float, item)) for item in [term["exponents"]] * len(term["coefficients"])]
        )
        coefficients.append([list(map(float, item)) for item in term["coefficients"]])

    return {
        "orbitals": list(itertools.chain(*orbitals)),
        "exponents": list(itertools.chain(*exponents)),
        "coefficients": list(itertools.chain(*coefficients)),
    }
